import torch


class Classifier(torch.nn.Module):
    def __init__(self, in_feats: int, num_cls: int) -> None:
        super().__init__()
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(in_features=in_feats, out_features=1024),
            torch.nn.ReLU(inplace=True),
            torch.nn.Linear(in_features=1024, out_features=512),
            torch.nn.ReLU(inplace=True),
            torch.nn.Linear(in_features=512, out_features=128),
            torch.nn.Dropout(inplace=True),
            torch.nn.ReLU(inplace=True),
            torch.nn.Linear(in_features=128, out_features=num_cls),
            torch.nn.Softmax(dim=1)
        )

    def forward(self, x):
        return self.fc(x)
